import pandas as pd
import numpy as np
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split
import torch

def load_data(file_path):
    """
    加载CSV文件数据
    
    参数:
    file_path (str): CSV文件的路径
    
    返回:
    pandas.DataFrame: 加载的数据
    """
    return pd.read_csv(file_path)

def preprocess_data(df):
    """
    预处理数据
    
    参数:
    df (pandas.DataFrame): 原始数据框
    
    返回:
    pandas.DataFrame: 预处理后的数据框
    """
    # 将时间戳转换为datetime类型
    df['timestamp'] = pd.to_datetime(df['timestamp'], unit='s')
    
    # 按船舶ID和时间戳排序
    df = df.sort_values(['mmsi', 'timestamp'])
    
    # 提取时间特征
    df['hour'] = df['timestamp'].dt.hour
    df['day'] = df['timestamp'].dt.day
    df['month'] = df['timestamp'].dt.month
    
    return df

def create_sequences(data, seq_length):
    """
    创建序列数据
    
    参数:
    data (numpy.array): 输入数据
    seq_length (int): 序列长度
    
    返回:
    tuple: (X, y) 其中X是输入序列，y是目标值
    """
    X, y = [], []
    for i in range(len(data) - seq_length):
        X.append(data[i:(i + seq_length)])
        y.append(data[i + seq_length, :2])  # 只取前两列 (lat, lon)
    return torch.FloatTensor(X), torch.FloatTensor(y)

def prepare_data(df, seq_length=10):
    """
    准备模型输入数据
    
    参数:
    df (pandas.DataFrame): 预处理后的数据框
    seq_length (int): 序列长度，默认为10
    
    返回:
    tuple: (X_train, X_test, y_train, y_test, scaler)
    """
    # 选择特征
    features = ['lat', 'lon', 'hour', 'day', 'month']
    
    # 使用MinMaxScaler进行特征缩放
    scaler = MinMaxScaler()
    scaled_data = scaler.fit_transform(df[features])
    
    # 创建序列数据
    X, y = create_sequences(scaled_data, seq_length)
    
    # 分割训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    
    return X_train, X_test, y_train, y_test, scaler